package com.zx.sa.service;

import com.alibaba.fastjson.JSONObject;
import com.google.gson.Gson;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.common.IndexParam.MetricType;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.DescribeCollectionReq;
import io.milvus.v2.service.collection.request.DropCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import io.milvus.v2.service.collection.request.LoadCollectionReq;
import io.milvus.v2.service.collection.request.ReleaseCollectionReq;
import io.milvus.v2.service.collection.response.DescribeCollectionResp;
import io.milvus.v2.service.collection.response.ListCollectionsResp;
import io.milvus.v2.service.partition.request.ListPartitionsReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.response.InsertResp;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

@Service
@Slf4j
public class MilvusService {

    public static final String COLLECTION_NAME = "ai_vectors";
    public static final int VECTOR_DIM = 128;

    @Autowired
    private MilvusClientV2 client;

    /**
     * 创建集合
     */
    public void createCollection() {
        // 检查集合是否存在
        if (client.hasCollection(HasCollectionReq.builder().collectionName(COLLECTION_NAME).build())) {
            log.info("集合已经存在, {}", COLLECTION_NAME);
            return;
        }

        // 定义集合结构
        CreateCollectionReq.CollectionSchema schema = client.createSchema();
        schema.addField(AddFieldReq.builder()
                .fieldName("my_id")
                .dataType(DataType.Int64)
                .isPrimaryKey(true)
                .autoID(false)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("my_vector")
                .dataType(DataType.FloatVector)
                .dimension(5)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("my_varchar")
                .dataType(DataType.VarChar)
                .maxLength(512)
                .build());

        // 添加索引[可选]
        IndexParam indexParamForIdField = IndexParam.builder()
                .fieldName("my_id")
                .indexType(IndexParam.IndexType.AUTOINDEX)
                .build();

        IndexParam indexParamForVectorField = IndexParam.builder()
                .fieldName("my_vector")
                .indexType(IndexParam.IndexType.AUTOINDEX)
                .metricType(MetricType.COSINE)  //余弦相似度 (COSINE)
                .build();

        List<IndexParam> indexParams = new ArrayList<>();
        indexParams.add(indexParamForIdField);
        indexParams.add(indexParamForVectorField);

        // 开始创建集合
        CreateCollectionReq customizedSetupReq = CreateCollectionReq.builder()
                .collectionName(COLLECTION_NAME)
                .collectionSchema(schema)
                .indexParams(indexParams)
                // 【可选】启用 mmap，这样可以减少内存占用，提高 Collections 的容量
                //.property(Constant.MMAP_ENABLED, "false")
                // 【可选】设置 Collections TTL，单位秒，一天后自动删除
                //.property(Constant.TTL_SECONDS, "86400")
                .build();
        client.createCollection(customizedSetupReq);
        log.info("创建集合, {}", COLLECTION_NAME);
    }

    public void collectionInfo() {
        ListCollectionsResp resp = client.listCollections();
        log.info("查看集合列表, {}", resp.getCollectionNames());

        for (String collectionName : resp.getCollectionNames()) {
            DescribeCollectionReq request = DescribeCollectionReq.builder()
                    .collectionName(COLLECTION_NAME)
                    .build();
            DescribeCollectionResp describeCollection = client.describeCollection(request);
            log.info("集合描述, collectionName={}, describe={}", collectionName, describeCollection.toString());

            ListPartitionsReq listPartitionsReq = ListPartitionsReq.builder()
                    .collectionName(COLLECTION_NAME)
                    .build();
            List<String> partitionNames = client.listPartitions(listPartitionsReq);
            log.info("集合分区信息, collectionName={}, partitionNames={}", collectionName, partitionNames);
        }
    }

    public void loadCollection() {
        // 加载集合到内存，提升查询速度
        LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                .collectionName(COLLECTION_NAME)
                // 【可选】可以加载整个集合，也可以只加载集合中的字段
                //.loadFields(Arrays.asList("my_id", "my_vector"))
                .build();
        client.loadCollection(loadCollectionReq);

        // 查看load状态
        GetLoadStateReq customSetupLoadStateReq = GetLoadStateReq.builder()
                .collectionName(COLLECTION_NAME)
                .build();
        Boolean loaded = client.getLoadState(customSetupLoadStateReq);
        log.info("加载集合到内存，提升查询速度, {}={}", COLLECTION_NAME, loaded);
    }

    public void releaseCollection() {
        ReleaseCollectionReq releaseCollectionReq = ReleaseCollectionReq.builder()
                .collectionName(COLLECTION_NAME)
                .build();

        client.releaseCollection(releaseCollectionReq);

        GetLoadStateReq loadStateReq = GetLoadStateReq.builder()
                .collectionName(COLLECTION_NAME)
                .build();
        Boolean res = client.getLoadState(loadStateReq);
        log.info("释放加载的集合内存, {}={}", COLLECTION_NAME, res);
    }

    public void insertData() {
        Gson gson = new Gson();
        List<JSONObject> data = Arrays.asList(
                JSONObject.parseObject(
                        "{\"my_id\": 0, \"my_vector\": [0.3580376395471989f, -0.6023495712049978f, 0.18414012509913835f, -0.26286205330961354f, 0.9029438446296592f], \"my_varchar\": \"pink_8682\"}"),
                JSONObject.parseObject(
                        "{\"my_id\": 1, \"my_vector\": [0.19886812562848388f, 0.06023560599112088f, 0.6976963061752597f, 0.2614474506242501f, 0.838729485096104f], \"my_varchar\": \"red_7025\"}"),
                JSONObject.parseObject(
                        "{\"my_id\": 2, \"my_vector\": [0.43742130801983836f, -0.5597502546264526f, 0.6457887650909682f, 0.7894058910881185f, 0.20785793220625592f], \"my_varchar\": \"orange_6781\"}"
                ),
                JSONObject.parseObject(
                        "{\"my_id\": 3, \"my_vector\": [0.3172005263489739f, 0.9719044792798428f, -0.36981146090600725f, -0.4860894583077995f, 0.95791889146345f], \"my_varchar\": \"pink_9298\"}"
                ),
                JSONObject.parseObject(
                        "{\"my_id\": 4, \"my_vector\": [0.4452349528804562f, -0.8757026943054742f, 0.8220779437047674f, 0.46406290649483184f, 0.30337481143159106f], \"my_varchar\": \"red_4794\"}"
                ),
                JSONObject.parseObject(
                        "{\"my_id\": 5, \"my_vector\": [0.985825131989184f, -0.8144651566660419f, 0.6299267002202009f, 0.1206906911183383f, -0.1446277761879955f], \"my_varchar\": \"yellow_4222\"}"
                ),
                JSONObject.parseObject(
                        "{\"my_id\": 6, \"my_vector\": [0.8371977790571115f, -0.015764369584852833f, -0.31062937026679327f, -0.562666951622192f, -0.8984947637863987f], \"my_varchar\": \"red_9392\"}"
                ),
                JSONObject.parseObject(
                        "{\"my_id\": 7, \"my_vector\": [-0.33445148015177995f, -0.2567135004164067f, 0.8987539745369246f, 0.9402995886420709f, 0.5378064918413052f], \"my_varchar\": \"grey_8510\"}"
                ),
                JSONObject.parseObject(
                        "{\"my_id\": 8, \"my_vector\": [0.39524717779832685f, 0.4000257286739164f, -0.5890507376891594f, -0.8650502298996872f, -0.6140360785406336f], \"my_varchar\": \"white_9381\"}"
                ),
                JSONObject.parseObject(
                        "{\"my_id\": 9, \"my_vector\": [0.5718280481994695f, 0.24070317428066512f, -0.3737913482606834f, -0.06726932177492717f, -0.6980531615588608f], \"my_varchar\": \"purple_4976\"}")
        );

        InsertReq insertReq = InsertReq.builder()
                .collectionName(COLLECTION_NAME)
                //选择插入的分区
                //.partitionName("partitionA")
                .data(data)
                .build();

        InsertResp insertResp = client.insert(insertReq);
        log.info("插入数据， {}", insertResp);
    }


    public void removeCollection() {
        DropCollectionReq dropQuickSetupParam = DropCollectionReq.builder()
                .collectionName(COLLECTION_NAME)
                .build();

        client.dropCollection(dropQuickSetupParam);
    }

}